import sys
import argparse
import uuid
from openai import OpenAI
from rich.console import Console
from prompt_toolkit import print_formatted_text
from prompt_toolkit.formatted_text import FormattedText
import pyfiglet
import time
from .agents.pentest_agent import PentestAgent
from .agents.chat_agent import ChatAgent
from .utils.web_login import login_via_web_app
import os
import json
import requests
from rich.console import Console

# Directory to store the data
DATA_DIR = os.path.join(os.path.dirname(__file__), '../logs/')

class PentestMuse:
    """
    Main class for Pentest Muse application.
    Handles the overall workflow and interactions between different agents.
    """
    def __init__(self):
        self.console = Console()
        self.pentest_agent = None
        self.command_agent = None

    def print_banner(self):
        self.console.clear()
        banner = pyfiglet.figlet_format("Pentest Muse", font="slant")
        self.console.print(banner, style="bold blue")
        time.sleep(0.5)

    def init_credentials(self):
        self.backend_mode = 'cli'
        self.credentials = None
        credentials_path = os.path.join(os.path.dirname(__file__), '../credentials.json')
        if os.path.exists(credentials_path):
            with open(credentials_path, 'r') as file:
                credentials_file = json.load(file)
                self.backend_mode = credentials_file.get('backend_mode')
                if self.backend_mode == 'openai':
                    self.credentials = credentials_file.get('openai_api_key')
                elif self.backend_mode == 'cli':
                    self.credentials = credentials_file.get('token')

        # reset openai api key if inputted as argument
        if self.args.openai_api_key:
            self.credentials = self.args.openai_api_key
            self.backend_mode = 'openai'
            # Write the token to the JSON file
            with open(credentials_path, 'w') as file:
                json.dump({'openai_api_key': self.credentials, 'backend_mode': 'openai'}, file)

        if self.args.login:
            # reset token if login is requested
            token = login_via_web_app(greeting=False)
            with open(credentials_path, 'w') as file:
                json.dump({'token': token, 'backend_mode': 'cli'}, file)
        elif self.backend_mode == 'openai':
            # test if openai api key is valid
            try:
                # ipdb.set_trace()
                test_client = OpenAI(api_key=self.credentials)
                test_client.chat.completions.create(
                    model="gpt-4-1106-preview",
                    messages=[{"role": "user", "content": 'Hello'}],
                    stream=True
                )
            except Exception as e:
                print_formatted_text(FormattedText([("bold red", f"\nSomething went wrong with your OpenAI API key. Please restart the application with a valid key or login via the web app.")]))
                self.credentials = login_via_web_app(greeting=False)
                self.backend_mode = 'cli'
                with open(credentials_path, 'w') as file:
                    json.dump({'token': self.credentials, 'backend_mode': 'cli'}, file)
        elif self.backend_mode == 'cli':
            # test if token is valid
            if self.credentials is not None:
                # test if token is valid
                try: 
                    response = requests.get('https://api.pentestmuse.ai/verifyToken', headers={'Authorization': f'Bearer {self.credentials}'})
                    if response.status_code == 200:
                        pass
                    else:
                        print_formatted_text(FormattedText([("bold yellow", f"\n Your credential is expired, please login again.")]))
                        self.credentials = login_via_web_app(greeting=False)
                        with open('credentials.json', 'w') as file:
                            json.dump({'token': self.credentials, 'backend_mode': 'cli'}, file)
                except requests.exceptions.ConnectionError:
                    # say that something is wrong with the connection
                    print_formatted_text(FormattedText([("bold red", f"\nSomething went wrong while connecting to the server. Please try again later.")]))
                    sys.exit()
                except Exception as e:
                    print_formatted_text(FormattedText([("bold red", f"\nSomething went wrong with your token. Please restart the application with a valid token or login via the web app.")]))
                    sys.exit()
            else:
                self.credentials = login_via_web_app(greeting=True)
                with open(credentials_path, 'w') as file:
                    json.dump({'token': self.credentials, 'backend_mode': 'cli'}, file)


    def run(self):
        self.print_banner()
        parser = argparse.ArgumentParser(description="Pentest Muse Application")
        parser.add_argument('mode', nargs='?', default='chat', help="Start mode of the application: 'chat' (default) or 'agent'")
        parser.add_argument('--openai-api-key', help="OpenAI API key. If not provided, login via web app is assumed.")
        parser.add_argument('--login', action='store_true', help="Login via web app.")
        self.args = parser.parse_args()

        # Initialize credentials
        self.init_credentials()

        # Initialize the language model
        if self.backend_mode == 'cli':
            from .language_models.pentest_muse_language_model import PentestMuseLanguageModel
            self.client = PentestMuseLanguageModel(credentials=self.credentials)
        elif self.backend_mode == 'openai':
            from .language_models.openai_language_model import OpenAILanguageModel
            self.client = OpenAILanguageModel(credentials=self.credentials)
        else:
            raise Exception("Invalid backend mode. Please use 'cli' or 'openai'.")

        # Choose between chat mode or action mode
        if self.args.mode == 'agent':
            task_id = str(uuid.uuid4())
            self.pentest_agent = PentestAgent(client=self.client, data_dir=DATA_DIR, task_id=task_id)
            self.pentest_agent.start()
        else:    
            task_id = str(uuid.uuid4())
            self.chat_agent = ChatAgent(client=self.client, data_dir=DATA_DIR, task_id=task_id)
            self.chat_agent.start()


def main():
    app = PentestMuse()
    app.run()

if __name__ == "__main__":
    main()